/*
 * Copyright 2019 WeBank
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.webank.wedatasphere.linkis.resourcemanager.utils

import java.security.PrivilegedExceptionAction
import java.util

import com.webank.wedatasphere.linkis.common.utils.{Logging, Utils}
import com.webank.wedatasphere.linkis.hadoop.common.utils.HDFSUtils
import com.webank.wedatasphere.linkis.resourcemanager.YarnResource
import com.webank.wedatasphere.linkis.resourcemanager.exception.RMErrorException
import org.apache.hadoop.yarn.api.records.{ApplicationReport, QueueInfo, YarnApplicationState}
import org.apache.hadoop.yarn.client.api.YarnClient
import org.apache.hadoop.yarn.conf.YarnConfiguration

import scala.collection.JavaConversions._

/**
  * Created by shanhuang on 2018/9/24.
  */
object YarnUtil extends Logging {

  private var yarnConf: YarnConfiguration = _

  private var yarnClient: YarnClient = _

  private def init(): Unit = {
    yarnConf = new YarnConfiguration()
    HDFSUtils.fillConfiguration(yarnConf)
    Utils.tryAndErrorMsg({
      HDFSUtils.getUserGroupInformation.doAs(new PrivilegedExceptionAction[Unit] {
        def run: Unit = {
          yarnClient = YarnClient.createYarnClient()
          yarnClient.init(yarnConf)
          yarnClient.start()
          logger.info("Yarn 客户端初始化完成")
        }
      })
    })("Failed to get YarnClient")
  }

  init()


  def getQueueInfo(queueName: String): (YarnResource, YarnResource) = {
    Utils.tryCatch {
      val info = HDFSUtils.getUserGroupInformation.doAs(new PrivilegedExceptionAction[QueueInfo] {
        override def run(): QueueInfo = {
          yarnClient.getQueueInfo(queueName)
        }
      })
      val queueStatistics = info.getQueueStatistics
      //总内存
      val totalMemory = (queueStatistics.getAllocatedMemoryMB + queueStatistics.getAvailableMemoryMB + queueStatistics.getPendingMemoryMB + queueStatistics.getReservedMemoryMB) * 1024l * 1024l
      val usedMemeory = totalMemory - queueStatistics.getAvailableMemoryMB * 1024l * 1024l
      //总核数
      val totalCore = (queueStatistics.getAllocatedVCores + queueStatistics.getAvailableVCores + queueStatistics.getPendingVCores + queueStatistics.getReservedVCores).toInt
      val usedCore = totalCore - queueStatistics.getAvailableVCores.toInt
      (new YarnResource(totalMemory, totalCore, 0, queueName), new YarnResource(usedMemeory, usedCore, 0, queueName))
    }(t => throw new RMErrorException(11006, "Get the Yarn queue information exception" + ".(获取Yarn队列信息异常)", t))
  }

  def getApplicationsInfo(queueName: String): Array[YarnAppInfo] = {
    Utils.tryCatch{
      val applicationReports = HDFSUtils.getUserGroupInformation.doAs(new PrivilegedExceptionAction[util.List[ApplicationReport]] {
        override def run(): util.List[ApplicationReport] = {
          yarnClient.getApplications(setAsJavaSet(Set(queueName)),null,null, util.EnumSet.range(YarnApplicationState.NEW,YarnApplicationState.RUNNING))
        }
      })
      applicationReports.map(app => {
        val usageReport = app.getApplicationResourceUsageReport
        val resources = usageReport.getUsedResources
        YarnAppInfo(app.getApplicationId.toString,app.getUser,app.getApplicationType,app.getYarnApplicationState.name(),new YarnResource(resources.getMemorySize*1024L*1024L,resources.getVirtualCores,0,app.getQueue))
      }).toArray
    }( t => throw new RMErrorException(11006, "Get the Yarn Application information exception.(获取Yarn Application信息异常)", t))
  }

  def getAllQueues: Array[String] = Utils.tryCatch {
    val queues = HDFSUtils.getUserGroupInformation.doAs(new PrivilegedExceptionAction[List[QueueInfo]] {
      override def run():List[QueueInfo] = yarnClient.getAllQueues.toList
    })
    queues.map(f => f.getQueueName).toArray
  } {
    t: Throwable => {
      logger.warn("get own queues error, msg: {}", t)
      Array()
    }
  }
}
